import backtrader as bt
import pandas as pd
import datetime
from strategy.utils.BaoStockPandasData import BaoStockPandasData

from utils.DataSource import get_stock_data2

class MACDStrategy(bt.Strategy):
    params = (
        ('macd_fast', 12),  # MACD快速EMA的周期
        ('macd_slow', 26),  # MACD慢速EMA的周期
        ('macd_signal', 9),  # MACD信号线的周期
        ('order_percentage', 0.95),  # 交易资金占总资金的百分比
    )

    def __init__(self):
        self.data_close = self.datas[0].close
        self.order = None
        self.macd = bt.indicators.MACDHisto(self.data_close,
                                           period_me1=self.params.macd_fast,
                                           period_me2=self.params.macd_slow,
                                           period_signal=self.params.macd_signal)
        
    def next(self):
        if self.order:
            return  # 如果有订单在执行，则不执行新的买卖操作

        # 当MACD柱状图由负转正，即MACD线从下穿信号线转为上穿，视为买入信号
        if self.macd.histo[-1] < 0 < self.macd.histo[0] and not self.position:
            amount_to_invest = (self.params.order_percentage * self.broker.cash) / self.data_close[0]
            self.buy(size=amount_to_invest)
            self.log(f"BUY EXECUTED --- Price: {self.data_close[0]}")

        # 当MACD柱状图由正转负，即MACD线从上穿信号线转为下穿，视为卖出信号
        elif self.macd.histo[-1] > 0 > self.macd.histo[0] and self.position:
            self.sell()
            self.log(f"SELL EXECUTED --- Price: {self.data_close[0]}")

    def log(self, txt, dt=None):
        ''' Logging function for this strategy'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return  # 订单状态未最终确定，等待

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f"BUY EXECUTED --- Price: {order.executed.price}, Cost: {order.executed.value}, Commission: {order.executed.comm}")
            else:  # Sell
                self.log(f"SELL EXECUTED --- Price: {order.executed.price}, Size: {order.executed.size}, Commission: {order.executed.comm}")

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        self.order = None

# 回测设置示例
if __name__ == '__main__':
    cerebro = bt.Cerebro()

    # 假设已经通过某种方式（如Yahoo Finance）获取了数据
    #data = bt.feeds.YahooFinanceData(dataname='AAPL', fromdate=datetime(2020, 1, 1), todate=datetime(2021, 12, 31))
    df = get_stock_data2("sh.603057", "2020-01-01", "2024-09-27")
    data = BaoStockPandasData(dataname=df)
    cerebro.adddata(data)

    cerebro.addstrategy(MACDStrategy, macd_fast=12, macd_slow=26, macd_signal=9)

    cerebro.broker.setcash(100000.0)
    cerebro.broker.set_coc(True)  # 现金交易模式

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    cerebro.run()
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    
    #在Backtrader中实现MACD（Moving Average Convergence Divergence）策略，主要是通过计算MACD线（快速EMA和慢速EMA的差值）、信号线（MACD线的EMA）以及MACD柱状图（MACD线与信号线之差），并根据它们之间的交叉或特定阈值来生成交易信号。以下是一个简单的MACD策略示例：
    #在这个策略中，我们直接使用了Backtrader内置的bt.indicators.MACDHisto指标来计算MACD线、信号线及柱状图，并基于MACD柱状图的正负变化来生成买入或卖出信号。用户可以根据自己的需求调整快速EMA、慢速EMA和信号线的周期参数。